library(monocle3)
library(tidyverse)
library(tidySingleCellExperiment)

# example data from C. elegans
cell_metadata <- 
  system.file('extdata',
              'worm_embryo/worm_embryo_coldata.rds',
              package='monocle3') |>
  readRDS()

gene_metadata <- 
  system.file('extdata',
              'worm_embryo/worm_embryo_rowdata.rds',
              package='monocle3') |>
  readRDS()

expression_matrix <-
  system.file('extdata',
              'worm_embryo/worm_embryo_expression_matrix.rds',
              package='monocle3') |>
  readRDS()
expression_matrix[1:5,1:5]

# create sce-based cds object ----------
cds <- expression_matrix |>
  new_cell_data_set(cell_metadata=cell_metadata,
                    gene_metadata=gene_metadata)

# preprocess ----------
cds <- preprocess_cds(cds)

## remove batch effects by batchelor ------
## generate adjusted PC reducedDim: Aligned
cds <- cds |> align_cds(alignment_group = "batch")

## non-linear dimred ----------
## use UMAP by default. set umap.fast_sgd=TRUE to speed up in cost of small inconsistency
## set cores > 1 to speed up in cost of small inconsistency
cds <- reduce_dimension(cds)

## cluster the cells (default by leiden) --------
## default will auto determine resolution, res=1e-4, generate 4 clusters
## if set res=.8, generate 122 clusters
cds <- cluster_cells(cds, verbose = T)

## leiden cluster
plot_cells(cds, color_cells_by = 'cluster')

cds$leiden_cluster <- clusters(cds)

## PAGA partition (auto done by cluster_cells)
plot_cells(cds, color_cells_by = 'partition')

cds$paga_partition <- partitions(cds)

# find cluster markers --------
marker_test_res <- cds |>
  top_markers(group_cells_by="partition",
              cores = 4)

head(marker_test_res)

top_specific_markers <- marker_test_res |>
  filter(fraction_expressing >= 0.10) |>
  group_by(cell_group) |>
  top_n(3, pseudo_R2)

top_specific_marker_ids <- top_specific_markers |>
  pull(gene_id) |>
  unique()

# dot plot
cds |>
  plot_genes_by_group(top_specific_marker_ids,
                      group_cells_by="partition")

# interactively subset cells on umap
cds_subset <- choose_cells(cds)

# ref-base auto annotate by Garnett ---------
curated_marker <- cds |>
  top_markers(group_cells_by = 'cell.type',
              cores = 4)

# Require at least JS specificty score > 0.5
garnett_markers <- curated_marker %>%
  filter(marker_test_q_value < 0.01 & specificity >= 0.5) %>%
  group_by(cell_group) %>%
  top_n(5, marker_score)

# Exclude genes that are good markers for more than one cell type:
garnett_markers <- garnett_markers %>% 
  group_by(gene_short_name) %>%
  filter(n() == 1)

garnett_markers |>
  generate_garnett_marker_file("0_Vignettes/garnett_worm_markers.txt")

## Install the monocle3 branch of garnett
devtools::install_github("cole-trapnell-lab/garnett", ref="monocle3")

library(garnett)

worm_classifier <- cds |>
  train_cell_classifier(marker_file = "0_Vignettes/garnett_worm_markers.txt", 
                        db = org.Ce.eg.db::org.Ce.eg.db,
                        cds_gene_id_type = "ENSEMBL",
                        num_unknown = 50,
                        marker_file_gene_id_type = "SYMBOL",
                        cores=8)

cds_gr <- cds |>
  classify_cells(worm_classifier,
                 db = org.Ce.eg.db::org.Ce.eg.db,
                 cluster_extend = TRUE,
                 cds_gene_id_type = "ENSEMBL")

cds_gr |>
  plot_cells(group_cells_by="partition",
             color_cells_by="cluster_ext_type")
# pre-trained classifier can be found

# Learn a graph to fit trajec ----------
## principal curve
cds <- learn_graph(cds)

cds |>
  plot_cells(color_cells_by = 'cell.type',
             label_groups_by_cluster = FALSE,
             label_leaves = TRUE,
             label_branch_points = TRUE)

## Order cells
cds |>
  plot_cells(color_cells_by = 'embryo.time.bin',
             group_label_size = 3,
             label_groups_by_cluster = FALSE,
             label_leaves = TRUE,
             label_branch_points = TRUE)

## will open a shiny window to choose root node
cds <- order_cells(cds)

## subset branch interactively
cds_sub <- choose_graph_segments(cds)

## plot trajec on UMAP ------
plot_cells(cds, color_cells_by = 'cell.type')

## default umap to 2d, it's ok to plot 3d
cds <- cds |> reduce_dimension(max_components = 3)

plot_cells_3d(cds)

## pseudotime plot
plot_cells(cds, color_cells_by = 'pseudotime')

cds$pseudotime <- pseudotime(cds)

# DEG analysis ----------
# With regression: log(expr) = intercept + k * variable_interest
# fit to continuous variable
# 1 core take 54s to fit
# 4 cores ~ 40s
# 8 cores ~ 30s
# 16 cores ~ 43s
gene_fits <- cds |>
  fit_models(model_formula_str = "~embryo.time",
             cores = 4)

emb_time_terms <- gene_fits |>
  coefficient_table() |>
  filter(term == "embryo.time" & q_value < .05) |>
  select(gene_short_name, term, q_value, estimate, id)

emb_time_terms |>
  ggplot(aes(estimate, log(q_value))) +
  geom_point()

sig_gene <- emb_time_terms |>
  slice_min(q_value, n = 6) |>
  pull(id)

# must subset cds to plot violin
cds[rownames(cds) %in% sig_gene, ] |>
  plot_genes_violin(group_cells_by = 'embryo.time.bin',
                    ncol = 2)

## include batch var ----------
batched_fits <- cds |>
  fit_models('~embryo.time + batch', cores = 4) 

batched_deg <- batched_fits |>
  coefficient_table() |>
  filter(term == "embryo.time" & q_value < .05) |>
  select(gene_short_name, term, q_value, estimate, id)

# significant improvement can be observed after adding batch var
# no working with too many genes?
# no working with default quasi-possion model
compare_res <- batched_fits |>
  compare_models(gene_fits)

# With graph autocorrelation -----
# use knn graph by default
knn_res <- cds |>
  graph_test(cores = 8) |>
  filter(q_value < 0.05)

# more positive morans I indicate a gene expressed in a focal region of umap
knn_res |>
  ggplot(aes(morans_I, log(q_value))) +
  geom_point()

gene_module_df <- gene_module_df |>
  mutate(module = str_c('Module_', module))

# aggregate cells by clusters
cell_group_df <- cds@colData |>
  as_tibble() |>
  select(cell, cell.type)

agg_mat <- cds |>
  aggregate_gene_expression(gene_module_df,
                            cell_group_df)

agg_mat |> glimpse()

agg_mat |>
  pheatmap::pheatmap(angle_col = 45,
                     scale = 'column',
                     clustering_method = 'ward.D2')

## find modules of co-regulated genes ------------
# run UMAP and cluster on genes, res=.01 produce 34 clusters
gene_module_df <- cds[knn_res$id, ] |>
  find_gene_modules(resolution = .01, verbose = T)

# better use principal graph for trajec analysis -------
# need learn graph first
pr_test_res <- cds |>
  graph_test(neighbor_graph="principal_graph", cores=4)

pr_deg <- pr_test_res |>
  filter(q_value < 0.05) |>
  as_tibble()

morans_top <- pr_deg |>
  slice_max(morans_I, n=3) 

# plot genes along pseudotime -----------
cds[morans_top$id, ] |>
  plot_genes_in_pseudotime()

# project query onto ref -------
# Load the reference data set.
matrix_ref <- readMM(gzcon(url("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/cao.mouse_embryo.sample.mtx.gz")))
cell_ann_ref <- read.csv(gzcon(url("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/cao.mouse_embryo.sample.coldata.txt.gz"), text=TRUE), sep='\t')
gene_ann_ref <- read.csv(gzcon(url("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/cao.mouse_embryo.sample.rowdata.txt.gz"), text=TRUE), sep='\t')

cds_ref <- new_cell_data_set(matrix_ref,
                             cell_metadata = cell_ann_ref,
                             gene_metadata = gene_ann_ref)

# Load the query data set.
matrix_qry <-
url("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/srivatsan.mouse_embryo_scispace.sample.mtx.gz") |>
  gzcon() |>
  Matrix::readMM()

cell_ann_qry <-
tidyfst::fread("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/srivatsan.mouse_embryo_scispace.sample.coldata.txt.gz") |>
  column_to_rownames('V1')

gene_ann_qry <- 
  tidyfst::fread("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/srivatsan.mouse_embryo_scispace.sample.rowdata.txt.gz") |>
  column_to_rownames('V1')

cds_qry <- matrix_qry |>
  as('CsparseMatrix') |>
  new_cell_data_set(cell_metadata = cell_ann_qry,
                    gene_metadata = gene_ann_qry)
  
# Load the reference data set.
matrix_ref <-
  url("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/cao.mouse_embryo.sample.mtx.gz") |>
  gzcon() |>
  Matrix::readMM()

cell_ann_ref <- 
  tidyfst::fread("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/cao.mouse_embryo.sample.coldata.txt.gz") |>
  column_to_rownames('V1')

gene_ann_ref <- 
  tidyfst::fread("https://depts.washington.edu:/trapnell-lab/software/monocle3/mouse/data/cao.mouse_embryo.sample.rowdata.txt.gz") |>
  column_to_rownames('V1')

cds_ref <- matrix_ref |>
  as('CsparseMatrix') |>
  new_cell_data_set(cell_metadata = cell_ann_ref,
                    gene_metadata = gene_ann_ref)

genes_shared <- rownames(cds_ref) |>
  intersect(rownames(cds_qry))

# Remove non-shared genes.
cds_ref <- cds_ref[genes_shared,]
cds_qry <- cds_qry[genes_shared,]

cds_ref <- estimate_size_factors(cds_ref)
cds_qry <- estimate_size_factors(cds_qry)

# process the ref
cds_ref <- preprocess_cds(cds_ref)
cds_ref <- cds_ref |>
  reduce_dimension(build_nn_index = TRUE)

# Save the PCA, UMAP models & nn index for use with projection.
# a dir with 5 files
cds_ref |>
  save_transform_models('cds_ref_test_models')

# Load the reference transform models into the query cds.
cds_qry <- cds_qry |>
  load_transform_models('cds_ref_test_models')

# preprocess with transform keyword
cds_qry <- preprocess_transform(cds_qry)
cds_qry <- reduce_dimension_transform(cds_qry)

# Combine the reference and query data sets
cds_ref$data_set <- 'reference'
cds_qry$data_set <- 'query'

cds_combined <- list(cds_ref, cds_qry) |>
  combine_cds(cell_names_unique=TRUE, keep_reduced_dims=TRUE)
plot_cells(cds_combined, color_cells_by='data_set')

# transfer cell labels from ref
cds_qry_lab_xfr <- cds_qry |>
  transfer_cell_labels(ref_coldata=colData(cds_ref),
                       ref_column_name='Main_cell_type',
                       query_column_name='cell_type_xfr',
                       transform_models_dir='cds_ref_test_models')

cds_qry_lab_fix <- cds_qry_lab_xfr |>
  fix_missing_cell_labels(from_column_name='cell_type_xfr',
                          to_column_name='cell_type_fix')

# disk io of cds ----------
# special func io to dir to preserve UMAP model & NN index
cds_ref |>
  save_monocle_objects('0_Vignettes/cds_obj')

cds_ref <-
  load_monocle_objects('0_Vignettes/cds_obj')